package com.xh.hongaiagent.agent;


import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.xh.hongaiagent.agent.model.AgentStatus;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;

import java.util.List;
import java.util.stream.Collectors;

/**
 * 处理工具调用的基本代理类
 */
@EqualsAndHashCode(callSuper = true)
@Slf4j
@Data
public class ToolCallAgent extends ReActAgent{

    //可用的工具
    private final ToolCallback[] availableTools;

    //保存工具调用的结果
    private ChatResponse toolCallChatResponse;

    private final ToolCallbackProvider toolCallbackProvider;

    //工具调用管理类
    private final ToolCallingManager toolCallingManager;

    //工具调用的选项
    private ChatOptions toolCallOptions;

    public ToolCallAgent(ToolCallback[] availableTools, ToolCallbackProvider toolCallbackProvider) {
        super();
        this.availableTools = availableTools;
        this.toolCallbackProvider = toolCallbackProvider;
        this.toolCallingManager = ToolCallingManager.builder().build();
        //禁用SpringAI内置工具调用机制，使用自定义工具调用
        this.toolCallOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true)
                .build();
    }


    /**
     * 处理当前状态（context），决定下一步行动（是否需要调用工具，继续思考）
     * @return 下一步动作
     */
    @Override
    public boolean think() {
        List<Message> messages = getContext();
        //校验输入信息，拼接上一步骤的结果
        if(StrUtil.isNotBlank(getNextStatePrompt())){
            UserMessage userMessage = new UserMessage(getNextStatePrompt());
            messages.add(userMessage);
        }

        try {
            //调用大模型，获取模型输出
            Prompt prompt = new Prompt(messages,this.toolCallOptions);
            ChatResponse chatResponse = getChatClient()
                    .prompt(prompt)
                    .system(getSystemPrompt())
                    //工具调用
                    .tools(availableTools)
                    //mcp服务
                    .tools(toolCallbackProvider)
                    .call()
                    .chatResponse();
            this.toolCallChatResponse = chatResponse;
            //解析模型输出,获取调用工具的参数
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            List<AssistantMessage.ToolCall> toolCallList = assistantMessage.getToolCalls();
            String result = assistantMessage.getText();
            log.info(getName() + " thinking result: "+ result);
            //如果不需要调用工具
            if (toolCallList.isEmpty()){
                log.info(getName() + " does not need tools");
                getContext().add(assistantMessage);
                return false;
            }

            String callingresults = toolCallList.stream()
                    .map(toolCall -> String.format("Tool name: {}, Tool args:", toolCall.name(), toolCall.arguments()))
                    .collect(Collectors.joining("\n"));
            log.info(getName()+ "  choosed: "+ toolCallList.size()+ " tools to execute.");
            log.info("Thinking toolCallList including: \n"+ callingresults);
            return true;
        } catch (Exception e) {
            log.error(getName() + "  thinking occurred errors: "+ e.getMessage());
            getContext().add(new AssistantMessage("执行遇到错误: " + e.getMessage()));
            return false;
        }
        //处理异常
    }

    /**
     * 执行工具调用，并且返回调用结果
     * @return
     */
    @Override
    public String act() {
        if(!toolCallChatResponse.hasToolCalls()){
            return  "没有工具需要调用";
        }
        Prompt prompt = new Prompt(getContext(),this.toolCallOptions);
        //手动调用工具
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse);
        //将调用结果添加到上下文
        setContext(toolExecutionResult.conversationHistory());
        //获取最后一条消息，也就是工具调用结果
        ToolResponseMessage last = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
        //判断是否调用了终止工具
        boolean isTerminated = last.getResponses().stream()
                .anyMatch(response -> response.name().equals("doTerminate"));
        if (isTerminated){
            setStatus(AgentStatus.FINISHED);
        }

        String toolExecuteResult = last.getResponses().stream()
                .map(response -> "调用了 工具：" + response.name() + " 其执行结果为： " + response.responseData())
                .collect(Collectors.joining("\n"));
        log.info(toolExecuteResult);
        return toolExecuteResult;
    }
}
